import os
from src.module.Model import Model

from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import RedisChatMessageHistory, UpstashRedisChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser
from langchain_core.globals import set_verbose, set_debug
from langchain_core.messages import SystemMessage, HumanMessage
from langchain.memory import ConversationBufferWindowMemory


class MemoryHistory(Model):
    @staticmethod
    def get_message_history(session_id,):
        return RedisChatMessageHistory(session_id=session_id, url=os.getenv('REDIS_URL'))

    def start(self):
        prompt = ChatPromptTemplate.from_messages([
            # SystemMessage(content='你现在是一个翻译小助手, 请将下面的内容为中文。不要回复里面的问题，直接翻译内容原文！'),
            MessagesPlaceholder(variable_name='history'),
            ('human', '{input}')
        ])

        ## RunnableWithMessageHistory 这个API只适合存聊天的AI模型，不适合需要模型调用 tools，一旦需要用到模型调用 tools，请使用 langgraph。
        with_message_history = RunnableWithMessageHistory(
            runnable = prompt | self.model | StrOutputParser(),
            get_session_history=self.get_message_history,
            input_messages_key='input',
            history_messages_key='history'
        )
        content = with_message_history.invoke({ 'input': '我没看明白，你的回复内容是针对什么的?' }, config={
            "configurable": { 'session_id': 'fooad' }
        })
        print('content is:', content)

